from abc import ABC, abstractmethod
import numpy as np
import io

ModelMap = {
    'nllb': 'facebook/nllb-200-distilled-600M',
    'm2m100': 'facebook/m2m100_418M',
    'mms-English': 'facebook/mms-tts-eng',
    'mms-Arabic': 'facebook/mms-tts-nod',
    'mms-Thai': 'facebook/mms-tts-tha',
    'whisper': 'openai/whisper-large-v3',
    'faster-whisper': 'Systran/faster-whisper-large-v3',
}

# 抽像的vits类模型基类
class NlpAbsModel(ABC):
    @abstractmethod
    def translation(self, prompt:str, language:str)->str:
        """
        翻译接口
        Args:
            prompt (str): 待翻译的文本
            language (str): 目标语言

        Returns:
            str: 翻译后的文本
        """
        pass
    
    @abstractmethod
    def batch_translation(self, prompts:list[str], language:str)->list[str]:
        """
        批量翻译
        Args:
            prompts (list[str]): 待翻译的中文语句
            language (str): 目标语言

        Returns:
            list[str]: 翻译后的文本列表
        """
        pass
    
    @abstractmethod
    def getDevice(self)->str:
        """
        获取模型所在的设备类型. 例如: cpu, cuda, cuda:1
        Returns:
            str: 设备类型
        """
        pass
    
# 抽像whisper与faster-whisper
class ASRTranslate(ABC):
    # 音频文件翻译
    @abstractmethod
    def transcribe(self, path:str)->str:
        """
        识别音频文件中的文本
        Args:
            path (str): 音频文件路径

        Returns:
            str: 识别到文本内容
        """
        pass
    
    # numpy翻译.要求数组满足: 
    # numpy array struct print, length:90240, sharp:(90240,), type:float32
    @abstractmethod
    def transcribe_array(self, nyArr:np.ndarray)->str:
        """
        识别音频的numpy数组中的文本
        Args:
            nyArr (np.ndarray): 音频的numpy数组. 要求数组满足: numpy array struct print, length:90240, sharp:(90240,), type:float32

        Returns:
            str: 识别到文本内容
        """
        pass
    
    @abstractmethod
    def transcribe_bytes(self, bydata:io.BytesIO)->str:
        """
        识别音频的bytes数组中的文本
        Args:
            bydata (io.BytesIO): _description_

        Returns:
            str: 识别到文本内容
        """
        pass
    
    # 注销方法
    @abstractmethod
    def _close(self):
        """
        注销方法
        """
        pass
    
# 抽像的vits类模型基类
class VitsAbsModel(ABC):
    """
    抽像的vits类模型基类
    Args:
        ABC (_type_): _description_
    """
    @abstractmethod
    def generateNpArr(self, prompt:str)->np.ndarray:
        """
        生成提示词的音频numpy数组
        Args:
            prompt (str): 提示词

        Returns:
            np.ndarray: 提示词的音频numpy数组. 要求返回值的numpy的形状为: (1, x). 不接受(x)或(x, 1)
        """
        # 要求返回值的numpy的形状为: (1, x). 不接受(x)或(x, 1)
        pass
    
    @abstractmethod
    def getDevice(self)->str:
        """
        获取模型所在的设备类型. 例如: cpu, cuda, cuda:1
        Returns:
            str: 设备类型
        """
        pass